#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# extractWordsFromCELEX.py
# written by Tomonori Nagano <tnagano@gc.cuny.edu> 
# September 11, 2010

import nltk
import re
import random

from collections import defaultdict

corpus_root = nltk.data.find('corpora/celex2/english/')

##############################################
# NEED TO READ THE CELEX2 MANUAL TO USE THIS
# p.102 of manual
# COBUILD 17.9 million words
# 16,600,000 written and 1,300,000 spoken
#
# many different frequency versions: phonological freq, dict freq etc...
##############################################
def main():
	celex = defaultdict(dict)
	
	temp = [line.strip() for line in open(corpus_root+'efw/efw.cd','r').readlines()]
	word_frequency = [entry.split('\\') for entry in temp]
	# 1.IdNum 2.Word 3.IdNumLemma 4.Cob 5.CobDev 6.CobMln 7.CobLog 8.CobW
	# 9.CobWMln 10.CobWLog 11.CobS 12.CobSMln 13.CobSLog
	col_index = ('IdNum', 'Word', 'IdNumLemma', 'Cob', 'CobDev', 'CobMln', 'CobLog', 'CobW','CobWMln', 'CobWLog', 'CobS', 'CobSMln', 'CobSLog')
	for entry in word_frequency:
		for i in range(len(col_index)):
			celex[entry[0]][col_index[i]] = entry[i]
	
	
	temp = [line.strip() for line in open(corpus_root+'epw/epw.cd','r').readlines()]
	word_phonology = [entry.split('\\') for entry in temp]
	# 1.IdNum 2.Word 3.Cob 4.IdNumLemma 5.PronCnt 6.PronStatus 7.PhonStrsDISC
	# 8.PhonCVBr 9.PhonSylBCLX 10.PronStatus 11.PhonStrsDISC  12.PhonCVBr 13.PhonSylBCLX 
	col_index = ('IdNum', 'Word', 'Cob', 'IdNumLemma', 'PronCnt', 'PronStatus', 'PhonStrsDISC','PhonCVBr', 'PhonSylBCLX', 'PronStatus', 'PhonStrsDISC',  'PhonCVBr', 'PhonSylBCLX')
	for entry in word_phonology:
		for i in range(0,9):
			celex[entry[0]][col_index[i]] = entry[i]
		j = 1
		try: 
			while(len(entry)>9+j):
				celex[entry[0]][col_index[i]+str(j)] = entry[j+9]
				celex[entry[0]][col_index[i]+str(j)] = entry[j+10]
				celex[entry[0]][col_index[i]+str(j)] = entry[j+11]
				celex[entry[0]][col_index[i]+str(j)] = entry[j+12]
				j += 1
		except IndexError as e:
			pass
			
	temp = [line.strip() for line in open(corpus_root+'emw/emw.cd','r').readlines()]
	word_morphology = [entry.split('\\') for entry in temp]
	# 1.IdNum 2.Word 3.Cob 4.IdNumLemma 5.FlectType 6.TransInfl
	col_index = ('IdNum', 'Word', 'Cob', 'IdNumLemma', 'FlectType', 'TransInfl')
	for entry in word_morphology:
		for i in range(len(col_index)):
			celex[entry[0]][col_index[i]] = entry[i]
	
	# [r]:r, [l]:l
	# back vowels: U ('put'), u: ('boon'), aU ('brow'), O: ('born'), A: ('barn')
	# front vowels: i: ('bean'), I ('pit'), eI ('bay'), E ('pet'), & ('pat')
	output = []
	
	for i in range(1,len(celex)):
		entry = celex[str(i)]
		this_pronunciation = entry['PhonSylBCLX']
		# skip if fewer than 3 syllables
		if len(re.findall('\[',this_pronunciation)) < 3:
			pass
		# skip if the frequency of the word is zero
		elif not (int(entry['CobW']) > 0 or int(entry['CobS']) > 0):
			pass
		# there is more than three [l]
		elif len(re.findall('l',this_pronunciation)) > 2:
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['three_l', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])
		# there is more than three [r]
		elif len(re.findall('r',this_pronunciation)) > 2:
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['three_r', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])
		# if there are both [l] and [r]
		elif re.search('l',this_pronunciation) and re.search('r',this_pronunciation):
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['r_and_l', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])
		elif re.search('l[U|u|aU|O:|A:]',this_pronunciation):
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['l_backvowel', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])
		elif re.search('l[i:|I|eI|E|&]',this_pronunciation):
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['l_frontvowel', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])
		elif re.search('r[U|u|aU|O:|A:]',this_pronunciation):
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['r_backvowel', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])
		elif re.search('r[i:|I|eI|E|&]',this_pronunciation):
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['r_frontvowel', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])
		elif re.search('[U|u|aU|O:|A:]l',this_pronunciation):
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['backvowel_l', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])
		elif re.search('[i:|I|eI|E|&]l',this_pronunciation):
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['frontvowel_l', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])
		elif re.search('[U|u|aU|O:|A:]r',this_pronunciation):
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['backvowel_r', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])
		elif re.search('[i:|I|eI|E|&]r',this_pronunciation):
			this_pronunciation = convertCELEXToIPA(this_pronunciation)
			numSyllable = len(re.findall('\[',this_pronunciation))
			this_pronunciation2 = this_pronunciation.replace('[','')
			this_pronunciation2 = this_pronunciation2.replace(']','')
			output.append(['frontvowel_r', entry['Word'], numSyllable, this_pronunciation, this_pronunciation2, entry['FlectType'], entry['TransInfl'], entry['CobW'], entry['CobS'], i])

	types = ['three_l','three_r','r_and_l','l_backvowel','l_frontvowel','r_backvowel','r_frontvowel', \
		'backvowel_l','frontvowel_l','backvowel_r','frontvowel_r']

	output2 = []
	for type in types:
		thisList = [line for line in output if line[0] == type]
		if len(thisList) > 100:
			thisSampleList = random.sample(thisList,100)
		else:
			thisSampleList = thisList
		output2.extend(thisSampleList)
		
	random.shuffle(output2)
	
	header = ['type', 'word', 'numSyllable','syllabel', 'pronunciation', 'wordType', 'infl', 'writtenFreq', 'spokenFreq', 'celexNo']
	output2.insert(0,header)
	for j, entry in enumerate(output2):
		print "%7d\t" % (j),
		for k,item in enumerate(entry):
			print unicode(item).encode("utf-8"),"\t",	
		print ''
"""		print "%5d" % (j),
		for k,item in enumerate(entry):
			if k == 0:
				print "%13s" % (unicode(item).encode("utf-8")),
			elif k == 2:
				print "%40s" % (unicode(item).encode("utf-8")),
			elif k == 3:
				print "%40s" % item,
			elif k == 4 or k == 5 or k == 6 or k == 7 or k == 8:
				print "%8s" % (unicode(item).encode("utf-8")),
			else:
				print "%20s" % (unicode(item).encode("utf-8")),			
		print ''
"""

def convertCELEXToIPA(pronunciationKey):
	# changin celex phone transcript to LaTeX TIPA style
	# see http://www.let.uu.nl/~Hugo.Quene/personal/phonchar.html
	# also see the celex2 English manual 4-25 and 4-26
	pronunciationKey = pronunciationKey.replace('x',u'k')
	pronunciationKey = pronunciationKey.replace('tS',u'ʧ')
	pronunciationKey = pronunciationKey.replace('dZ',u'ʤ')
	pronunciationKey = pronunciationKey.replace('N',u'ŋ')
	pronunciationKey = pronunciationKey.replace('T',u'θ')
	pronunciationKey = pronunciationKey.replace('D',u'ð')
	pronunciationKey = pronunciationKey.replace('S',u'ʃ')
	pronunciationKey = pronunciationKey.replace('Z',u'ʒ')
	#
	pronunciationKey = pronunciationKey.replace('N,',u'n̩')
	pronunciationKey = pronunciationKey.replace('m,',u'm̩')
	pronunciationKey = pronunciationKey.replace('n,',u'n')
	pronunciationKey = pronunciationKey.replace('l,',u'l̩')
	pronunciationKey = pronunciationKey.replace('r*',u'ɚ')
	pronunciationKey = pronunciationKey.replace('j,',u'y')
	#
	pronunciationKey = pronunciationKey.replace('i:',u'i:')
	pronunciationKey = pronunciationKey.replace('A:',u'ɑ:')
	pronunciationKey = pronunciationKey.replace('O:',u'ɔ:')
	pronunciationKey = pronunciationKey.replace('u:',u'u:')
	pronunciationKey = pronunciationKey.replace('3:',u'ɚ')
	#
	pronunciationKey = pronunciationKey.replace('@r',u'ɚ')
	pronunciationKey = pronunciationKey.replace('@',u'ɘ')
	#
	pronunciationKey = pronunciationKey.replace('@U',u'oʊ')
	#
	pronunciationKey = pronunciationKey.replace('I',u'ɪ')
	pronunciationKey = pronunciationKey.replace('E',u'ɛ')
	pronunciationKey = pronunciationKey.replace('&',u'æ')
	pronunciationKey = pronunciationKey.replace('O',u'ɔ')
	pronunciationKey = pronunciationKey.replace('V',u'ʌ')
	pronunciationKey = pronunciationKey.replace('U',u'ʊ')
	#
	pronunciationKey = pronunciationKey.replace('&~',u'æ')
	pronunciationKey = pronunciationKey.replace('A~:',u'eɪ')
	pronunciationKey = pronunciationKey.replace('{~:',u'ɑn')
	pronunciationKey = pronunciationKey.replace('O:~',u'ɑn')
	#
	pronunciationKey = pronunciationKey.replace(' ','')
	#
	return pronunciationKey
	

main()

